/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved.
 * Global Trust Authority is licensed under the Mulan PSL v2.
 * You can use this software according to the terms and conditions of the Mulan PSL v2.
 * You may obtain a copy of Mulan PSL v2 at:
 *     http://license.coscl.org.cn/MulanPSL2
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY OR FIT FOR A PARTICULAR
 * PURPOSE.
 * See the Mulan PSL v2 for more details.
 */

use crate::config::{ASCEND_NPU_IMA_LOG, ASCEND_NPU_TCG_EVENT_LOG};
use base64::engine::general_purpose;
use base64::Engine;
use plugin_manager::PluginError;
use serde::{Deserialize, Deserializer, Serialize};
use std::collections::{HashMap};

#[derive(Debug, Serialize, Deserialize)]
pub struct AscendNpuEvidence {
    pub ak_cert: String,
    pub quote: Quote,
    pub pcrs: Pcrs,
    pub logs: Vec<Log>,
}

#[derive(Debug, Serialize, Deserialize)]
pub struct Quote {
    pub quote_data: String,
    pub signature: String,
}

#[derive(Debug, Serialize, Deserialize)]
pub struct Pcrs {
    pub hash_alg: String,
    pub pcr_values: Vec<PCR>,
}

#[derive(Debug, Serialize, Deserialize)]
pub struct PCR {
    pub pcr_index: u32,
    pub pcr_value: String,
}

#[derive(Debug, Serialize, Deserialize)]
pub struct Log {
    pub log_type: String,
    pub log_data: String,
}

#[derive(Debug, Serialize, Deserialize)]
pub struct DcmiEvidence {
    #[serde(with = "serde_bytes")]
    pub ak_cert: Vec<u8>,
    pub quote: DcmiQuote,
    #[serde(default, with = "serde_bytes")]
    pub boot_measurements: Vec<u8>,
    #[serde(default, with = "serde_bytes")]
    pub runtime_measurements: Vec<u8>,
}

#[derive(Debug, Serialize, Deserialize)]
pub struct DcmiQuote {
    #[serde(with = "serde_bytes")]
    pub data: Vec<u8>,
    #[serde(with = "serde_bytes")]
    pub signature: Vec<u8>,
    pub pcrs: DcmiPcrs,
}

#[derive(Debug, Serialize, Deserialize)]
pub struct DcmiPcrs {
    pub hash_alg: String,
    #[serde(deserialize_with = "deserialize_map")]
    pub values: HashMap<String, Vec<u8>>,
}

/// # deserialize_map
/// 
/// Deserializes a CBOR map into a `HashMap<Vec<u8>, Vec<u8>>`.
/// 
/// # Parameters
/// - `deserializer`: The deserializer instance.
/// 
/// # Errors
/// 
/// This function will return an error if:
/// - The CBOR value is not a map.
/// - The map keys or values are not CBOR byte strings.
/// 
/// # Returns
/// A `Result` containing the deserialized `HashMap<Vec<u8>, Vec<u8>>` or a `serde::de::Error`.
fn deserialize_map<'de, D>(deserializer: D) -> Result<HashMap<String, Vec<u8>>, D::Error>
where
    D: Deserializer<'de>,
{
    let raw_value = serde_cbor::Value::deserialize(deserializer)?;

    let cbor_map = match raw_value {
        serde_cbor::Value::Map(map) => map,
        _ => return Err(serde::de::Error::custom("Expected a CBOR Map")),
    };

    let mut result = HashMap::new();
    for (k, v) in cbor_map {
        if let (serde_cbor::Value::Text(k_str), serde_cbor::Value::Bytes(v_bytes)) = (k, v) {
            if !v_bytes.is_empty() && !k_str.is_empty() {
                result.insert(k_str, v_bytes);
            }
        }
    }
    Ok(result)
}

impl TryFrom<DcmiEvidence> for AscendNpuEvidence {
    type Error = PluginError;

    /// # try_from
    /// 
    /// Converts a `DcmiEvidence` instance into an `AscendNpuEvidence` instance.
    /// 
    /// # Parameters
    /// - `evidence`: The `DcmiEvidence` instance to convert.
    /// 
    /// # Errors
    /// 
    /// This function will return an error if:
    /// - The PCR values cannot be parsed as valid hexadecimal strings.
    /// - The log types are not valid.
    /// 
    /// # Returns
    /// A `Result` containing the `AscendNpuEvidence` instance or a `PluginError`.
    fn try_from(evidence: DcmiEvidence) -> Result<Self, PluginError> {
        let mut logs = Vec::new();
        if !evidence.boot_measurements.is_empty() {
            logs.push(Log {
                log_type: ASCEND_NPU_TCG_EVENT_LOG.to_string(),
                log_data: general_purpose::STANDARD.encode(evidence.boot_measurements),
            });
        }
        if !evidence.runtime_measurements.is_empty() {
            logs.push(Log {
                log_type: ASCEND_NPU_IMA_LOG.to_string(),
                log_data: general_purpose::STANDARD.encode(evidence.runtime_measurements),
            });
        }

        let mut pcrs = Pcrs {
            hash_alg: evidence.quote.pcrs.hash_alg,
            pcr_values: evidence
                .quote
                .pcrs
                .values
                .into_iter()
                .map(|(k, v)| {
                    let idx = k
                        .parse::<u32>()
                        .map_err(|e| PluginError::InputError(format!(
                            "Invalid PCR index '{}': {}",
                            k, e
                        )))?;

                    Ok(PCR {
                        pcr_index: idx,
                        pcr_value: hex::encode(v),
                    })
                })
                .collect::<Result<Vec<_>, PluginError>>()?,
        };
        pcrs.pcr_values.sort_by_key(|p| p.pcr_index);
        Ok(Self {
            ak_cert: general_purpose::STANDARD.encode(evidence.ak_cert),
            quote: Quote {
                quote_data: general_purpose::STANDARD.encode(evidence.quote.data),
                signature: general_purpose::STANDARD.encode(evidence.quote.signature),
            },
            pcrs,
            logs,
        })
    }
}
